
from flask import Flask, jsonify, abort, request
import time
import threading
from datetime import datetime, timedelta
import uuid
from seggpt import main as seggpt_main
import cv2

app = Flask(__name__)
result_dict = {}
status_dict = {}

@app.route('/')
def index():
    return 'service: seggpt-tracker'

@app.route('/meta')
def meta():
    return jsonify({
        'name': 'seggpt-tracker',
        'description': 'seggpt-tracker',
        'category': 'tracking',
    })

@app.route('/tasks')
def get_tasks():
    return jsonify([
        # 'main: FilePath -> List<Mask>', # not implemented
        'main_returnpath: FilePath -> FilePath',
        # 'test: Void -> Void' # not implemented
    ])

@app.route('/workflows')
def get_workflows():
    return jsonify({
        'workflows': []
    })

@app.route('/mount', methods=['POST'])
def mount():
    return str(uuid.uuid4())

def mark_as_failed(task_id: str, message: str):
    status_dict[task_id]['status'] = 'Failed'
    status_dict[task_id]['errmsg'] = message
    print(f'task {task_id} failed: {message}')

@app.route('/task/main_returnpath', methods=['POST'])
def main_returnpath_start_task():
    session_key = request.json.get('session_key')
    file_path = request.json.get('input_data')
    if not file_path:
        abort(400, 'file_path is required')
    if not isinstance(file_path, str):
        abort(400, 'file_path must be a string')
    global status_dict
    id = str(uuid.uuid4())
    status_dict[id] = {
        'status': 'Running',
        'percentage': 0
    }
    result_dict[id] = None
    def run_task(task_id: str):
        cap = cv2.VideoCapture(file_path)
        if not cap.isOpened():
            mark_as_failed(task_id, 'failed to open video')
        ret, frame0 = cap.read()
        if not ret:
            mark_as_failed(task_id, 'failed to read frame')
        try:
            mask_path = seggpt_main(cap, frame0, file_path, status_dict, task_id)
        except Exception as e:
            mark_as_failed(task_id, f'task {task_id} failed: {e}')
        if status_dict[task_id]['status'] == 'Stopped':
            print(f'task {task_id} aborted')
        if mask_path is None:
            result_dict[task_id] = None
            mark_as_failed(task_id, 'no result')
        else:
            result_dict[task_id] = mask_path
            status_dict[task_id]['status'] = 'Completed'
            status_dict[task_id]['percentage'] = 100

    task_thread = threading.Thread(target=run_task, args=(id,))
    task_thread.start()
    return jsonify({
        'task_id': id,
        'progress_pipe_addr': f'/task/main_returnpath/progress/{id}',
        'result_addr': f'/task/main_returnpath/result/{id}'
    })

@app.route('/task/main_returnpath/progress/<task_id>')
def main_returnpath_get_progress(task_id):
    if task_id not in status_dict:
        abort(400, 'task not found')
    return jsonify(status_dict[task_id])

@app.route('/task/main_returnpath/result/<task_id>')
def main_returnpath_get_result(task_id):
    if task_id not in result_dict:
        abort(400, 'task not found')
    return jsonify(result_dict[task_id])


@app.route('/task/main_returnpath/terminate/<task_id>', methods=['POST'])
def main_returnpath_terminate(task_id):
    if task_id not in status_dict:
        abort(400, 'task not found')
    status_dict[task_id]['status'] = 'Stopped'
    return jsonify(status_dict[task_id])

if __name__ == '__main__':
    app.run(host='127.0.0.1', port=10010, debug=True)
